// pc_udp_receiver.cpp — UDP 重组 -> Draco 解码 -> PointCloud2 发布
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud2_iterator.h>

#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#include <endian.h>

#include <unordered_map>
#include <vector>
#include <string>
#include <cstring>
#include <algorithm>
#include <chrono>
#include <thread>
#include <array>
#include <atomic>

// Draco
#include <draco/core/decoder_buffer.h>
#include <draco/compression/decode.h>
#include <draco/point_cloud/point_cloud.h>
#include <draco/attributes/geometry_attribute.h>

#pragma pack(push, 1)
struct UdpHeader {
  uint32_t magic;         // 'PCUD'
  uint32_t seq;
  uint16_t total_chunks;
  uint16_t idx;
  uint32_t payload_size;
  uint64_t stamp_ns;
  uint8_t  codec;         // 0=raw,1=zstd,2=draco
  uint8_t  reserved[3];
};
#pragma pack(pop)

struct FrameBuf {
  uint16_t total{0};
  std::vector<bool> got;
  std::vector<std::vector<uint8_t>> chunks;
  ros::Time stamp;
  std::chrono::steady_clock::time_point t_start;
  size_t bytes_acc{0};
};

class PcReceiverDraco {
public:
  explicit PcReceiverDraco(ros::NodeHandle& nh_priv) : nh_priv_(nh_priv) {
    nh_priv_.param<std::string>("out_cloud_topic", out_topic_, std::string("/rx_cloud_xyzrgb"));
    nh_priv_.param<std::string>("bind_ip", bind_ip_, std::string("0.0.0.0"));
    nh_priv_.param<int>("bind_port", bind_port_, 50000);
    nh_priv_.param<int>("rcvbuf_mb", rcvbuf_mb_, 32);
    nh_priv_.param<double>("reassembly_timeout_s", reasm_timeout_s_, 0.3);
    nh_priv_.param<std::string>("frame_id", frame_id_, std::string("map"));

    pub_ = nh_priv_.advertise<sensor_msgs::PointCloud2>(out_topic_, 1);

    sock_ = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock_ < 0) { perror("socket"); throw std::runtime_error("socket() failed"); }
    int rcvbuf = rcvbuf_mb_ * 1024 * 1024;
    setsockopt(sock_, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf));

    std::memset(&addr_, 0, sizeof(addr_));
    addr_.sin_family = AF_INET;
    addr_.sin_port   = htons(bind_port_);
    addr_.sin_addr.s_addr = (bind_ip_ == "0.0.0.0") ? INADDR_ANY : inet_addr(bind_ip_.c_str());

    if (bind(sock_, reinterpret_cast<sockaddr*>(&addr_), sizeof(addr_)) < 0) {
      perror("bind"); throw std::runtime_error("bind() failed");
    }

    ROS_INFO_STREAM("pc_udp_receiver_draco bind " << bind_ip_ << ":" << bind_port_
                    << " out=" << out_topic_ << " frame_id=" << frame_id_);

    recv_thread_ = std::thread([this](){ this->recvLoop(); });
  }

  ~PcReceiverDraco(){
    running_ = false;
    if (recv_thread_.joinable()) recv_thread_.join();
    if (sock_ >= 0) close(sock_);
  }

private:
  void recvLoop() {
    std::vector<uint8_t> buf(65536);
    while (running_ && ros::ok()) {
      sockaddr_in src{}; socklen_t slen = sizeof(src);
      ssize_t n = recvfrom(sock_, buf.data(), buf.size(), 0, reinterpret_cast<sockaddr*>(&src), &slen);
      if (n <= 0) continue;
      if ((size_t)n < sizeof(UdpHeader)) continue;

      UdpHeader h{};
      std::memcpy(&h, buf.data(), sizeof(UdpHeader));
      if (ntohl(h.magic) != 0x50435544) continue; // magic

      const uint32_t seq   = ntohl(h.seq);
      const uint16_t total = ntohs(h.total_chunks);
      const uint16_t idx   = ntohs(h.idx);
      const uint32_t sz    = ntohl(h.payload_size);
      const uint8_t  codec = h.codec;

      if (codec != 2) { // 只处理 Draco
        continue;
      }
      if (sizeof(UdpHeader) + sz != (size_t)n) continue;

      auto &fb = frames_[seq];
      if (fb.total == 0) {
        fb.total = total;
        fb.got.assign(total, false);
        fb.chunks.resize(total);
        fb.stamp.fromNSec(be64toh(h.stamp_ns));
        fb.t_start = std::chrono::steady_clock::now();
        fb.bytes_acc = 0;
      }
      if (idx < fb.total && !fb.got[idx]) {
        fb.chunks[idx].assign(buf.data()+sizeof(UdpHeader), buf.data()+sizeof(UdpHeader)+sz);
        fb.got[idx] = true;
        fb.bytes_acc += sz;
      }

      reassemblyGC();

      if (std::all_of(fb.got.begin(), fb.got.end(), [](bool v){return v;})) {
        std::vector<uint8_t> payload;
        payload.reserve(fb.bytes_acc);
        for (uint16_t i=0;i<fb.total;++i){
          payload.insert(payload.end(), fb.chunks[i].begin(), fb.chunks[i].end());
        }
        const ros::Time stamp = fb.stamp;
        frames_.erase(seq);
        publishDecoded(payload, stamp);
      }
    }
  }

  void reassemblyGC() {
    const auto now = std::chrono::steady_clock::now();
    for (auto it = frames_.begin(); it != frames_.end(); ) {
      const double dt = std::chrono::duration<double>(now - it->second.t_start).count();
      if (dt > reasm_timeout_s_) {
        it = frames_.erase(it);
      } else ++it;
    }
  }

  void publishDecoded(const std::vector<uint8_t>& payload, const ros::Time& stamp) {
    draco::DecoderBuffer dbuf;
    dbuf.Init(reinterpret_cast<const char*>(payload.data()), payload.size());
    draco::Decoder decoder;
    auto res = decoder.DecodePointCloudFromBuffer(&dbuf);
    if (!res.ok()) {
      ROS_WARN_THROTTLE(1.0, "Draco decode failed: %s", res.status().error_msg());
      return;
    }
    std::unique_ptr<draco::PointCloud> pc = std::move(res).value();
    const int N = pc->num_points();
    if (N <= 0) return;

    const draco::PointAttribute* pos_att =
        pc->GetNamedAttribute(draco::GeometryAttribute::POSITION);
    const draco::PointAttribute* col_att =
        pc->GetNamedAttribute(draco::GeometryAttribute::COLOR);

    sensor_msgs::PointCloud2 msg;
    msg.header.stamp = stamp;
    msg.header.frame_id = frame_id_;
    msg.height = 1;
    msg.width  = N;

    sensor_msgs::PointCloud2Modifier mod(msg);
    if (col_att) {
      mod.setPointCloud2FieldsByString(2, "xyz", "rgb");
    } else {
      mod.setPointCloud2FieldsByString(1, "xyz");
    }
    mod.resize(N);

    sensor_msgs::PointCloud2Iterator<float> it_x(msg, "x");
    sensor_msgs::PointCloud2Iterator<float> it_y(msg, "y");
    sensor_msgs::PointCloud2Iterator<float> it_z(msg, "z");
    std::unique_ptr<sensor_msgs::PointCloud2Iterator<uint8_t>> it_rgb;
    if (col_att) {
      it_rgb.reset(new sensor_msgs::PointCloud2Iterator<uint8_t>(msg, "rgb"));
    }

    draco::PointIndex pi(0);
    for (int i=0; i<N; ++i, ++it_x, ++it_y, ++it_z, pi = draco::PointIndex(i+1)) {
      if (pos_att) {
        const auto avi = pos_att->mapped_index(pi);
        std::array<float, 3> xyz{};
        pos_att->GetValue(avi, &xyz);
        *it_x = xyz[0]; *it_y = xyz[1]; *it_z = xyz[2];
      } else {
        *it_x = *it_y = *it_z = 0.f;
      }

      if (col_att && it_rgb) {
        const auto avi = col_att->mapped_index(pi);
        const int comps = col_att->num_components();
        if (comps >= 4) {
          std::array<uint8_t, 4> c4{};
          col_att->GetValue(avi, &c4);
          (*it_rgb)[0] = c4[0];
          (*it_rgb)[1] = c4[1];
          (*it_rgb)[2] = c4[2];
        } else if (comps == 3) {
          std::array<uint8_t, 3> c3{};
          col_att->GetValue(avi, &c3);
          (*it_rgb)[0] = c3[0];
          (*it_rgb)[1] = c3[1];
          (*it_rgb)[2] = c3[2];
        } else {
          (*it_rgb)[0] = 0;
          (*it_rgb)[1] = 0;
          (*it_rgb)[2] = 0;
        }
        ++(*it_rgb);
      }
    }

    pub_.publish(msg);
    ROS_INFO_THROTTLE(1.0, "[RX-Draco] pts=%d (rgb=%s) bytes=%zu",
                      N, col_att?"yes":"no", payload.size());
  }

  ros::NodeHandle nh_priv_;
  ros::Publisher pub_;
  std::string out_topic_, bind_ip_, frame_id_;
  int bind_port_{50000}, rcvbuf_mb_{32};
  double reasm_timeout_s_{0.3};

  int sock_{-1};
  sockaddr_in addr_{};
  std::thread recv_thread_;
  std::atomic<bool> running_{true};

  std::unordered_map<uint32_t, FrameBuf> frames_;
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "pc_udp_receiver_draco");
  ros::NodeHandle nh_priv("~");
  try { PcReceiverDraco r(nh_priv); ros::spin(); }
  catch (const std::exception& e) { ROS_FATAL("receiver init failed: %s", e.what()); return 1; }
  return 0;
}

